# m3e-base 模型演示
# 功能：文本向量生成、相似度计算、简单检索
import time

import chromadb
import numpy as np
from chromadb.config import Settings
from sentence_transformers import SentenceTransformer


def main():
    # 1. 加载m3e-base模型
    # 首次运行会自动下载模型，若网络较慢可手动下载后指定本地路径
    print("正在加载m3e-base模型...")

    model_dir = "AI-ModelScope/m3e-base"
    start_time = time.time()
    model = SentenceTransformer(model_dir)
    end_time = time.time()

    print(f"Model loaded in {end_time - start_time:.2f} seconds")
    # Our sentences we like to encode
    sentences = [
        '* Moka 此文本嵌入模型由 MokaAI 训练并开源，训练脚本使用 uniem',
        '* Massive 此文本嵌入模型通过**千万级**的中文句对数据集进行训练',
        '* Mixed 此文本嵌入模型支持中英双语的同质文本相似度计算，异质文本检索等功能，未来还会支持代码检索，ALL in one'
    ]

    # Sentences are encoded by calling model.encode()
    embeddings = model.encode(sentences)

    # Print the embeddings
    for sentence, embedding in zip(sentences, embeddings):
        print("Sentence:", sentence)
        print("Embedding:", embedding)
        print("")
    # 如需使用GPU加速，添加device参数：device="cuda"
    # model = SentenceTransformer("moka-ai/m3e-base", device="cuda")

    # 2. 文本向量生成演示
    print("\n=== 文本向量生成演示 ===")
    texts = [
        "苹果是一种常见的水果，味道酸甜",
        "香蕉富含钾元素，对心脏有益",
        "程序员经常使用Python语言开发人工智能应用",
        "机器学习是人工智能的一个重要分支",
        "西瓜是夏季很受欢迎的水果，水分充足"
    ]

    # 生成向量
    embeddings = model.encode(texts)

    # 输出向量信息
    print(f"生成的向量维度: {embeddings.shape[1]}")
    print(f"第一句文本的向量前5位: {embeddings[0][:5]}")

    # 3. 文本相似度计算
    print("\n=== 文本相似度计算 ===")
    # 计算第一句与其他句子的余弦相似度
    first_embedding = embeddings[0]
    for i, text in enumerate(texts):
        similarity = np.dot(first_embedding, embeddings[i]) / (
                    np.linalg.norm(first_embedding) * np.linalg.norm(embeddings[i]))
        print(f"与'苹果是一种常见的水果，味道酸甜'的相似度: {text} -> {similarity:.4f}")

    # 4. 简单检索演示（使用chromadb作为向量数据库）
    print("\n=== 简单检索演示 ===")
    # 初始化chromadb
    client = chromadb.Client(Settings(
        persist_directory="./chroma_db",  # 向量数据存储目录
        is_persistent=True
    ))

    # 创建集合
    collection = client.get_or_create_collection(name="demo_collection")

    # 向向量库中添加文档
    collection.add(
        documents=texts,
        embeddings=embeddings.tolist(),
        ids=[f"id_{i}" for i in range(len(texts))]
    )

    # 检索测试
    queries = [
        "哪些水果适合夏天吃？",
        "人工智能相关的编程语言有哪些？"
    ]

    for query in queries:
        # 生成查询向量
        query_embedding = model.encode([query])

        # 检索相似文档
        results = collection.query(
            query_embeddings=query_embedding.tolist(),
            n_results=2  # 返回最相似的2个结果
        )

        print(f"\n查询: {query}")
        print("检索结果:")
        for doc, score in zip(results['documents'][0], results['distances'][0]):
            # 注意：chromadb返回的是距离，值越小越相似
            print(f"  文档: {doc} (相似度: {1 - score:.4f})")


if __name__ == "__main__":
    main()
